#! /usr/bin/python3
# -*- coding: UTF-8 -*-
  
'''※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※
File Name: affine_cipher.py
Author: GID5564
Description: 仿射加密
Version: 1.0
Created Time: 08/12/24-14:23:10
※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※'''


import random
import sys

import cryptomath_module as cryptomath


# 定义符号集，假设为一个字符串常量
SYMBOLS = (
    r""" !"#$%&'()*+,-./0123456789:;<=>?@ABCDEFGHIJKLMNOPQRSTUVWXYZ[\]^_`"""
    r"""abcdefghijklmnopqrstuvwxyz{|}~"""
)




def check_keys(keyA, keyB, mode):
    """
    检查给定的密钥是否适用于仿射密码算法。

    参数:
    keyA (int): 仿射密码中的乘法密钥。
    keyB (int): 仿射密码中的加法密钥。
    mode (str): 操作模式，可以是 'encrypt' 或 'decrypt'。

    返回值:
    无返回值，如果密钥无效则程序会退出并显示错误信息。
    """
    
    # 如果模式是加密模式
    if mode == "encrypt":
        # 检查 keyA 是否为 1，如果是则退出并提示用户选择不同的密钥
        if keyA == 1:
            sys.exit(
                "当 key A 设置为 1 时，仿射密码变得很弱。请选择不同的密钥。"
            )
        # 检查 keyB 是否为 0，如果是则退出并提示用户选择不同的密钥
        if keyB == 0:
            sys.exit(
                "当 key B 设置为 0 时，仿射密码变得很弱。请选择不同的密钥。"
            )
    
    # 检查 keyA 和 keyB 是否在有效范围内
    if keyA < 0 or keyB < 0 or keyB > len(SYMBOLS) - 1:
        sys.exit(
            f"key A 必须大于 0，key B 必须在 0 到 {len(SYMBOLS) - 1} 之间。"
        )
    
    # 检查 keyA 和符号集大小是否互质
    if cryptomath.gcd(keyA, len(SYMBOLS)) != 1:
        sys.exit(
            f"key A {keyA} 和符号集大小 {len(SYMBOLS)} 不是互质的。请选择不同的密钥。"
        )


def encrypt_message(key: int, message: str) -> str:
    """
    使用仿射密码对消息进行加密。

    参数:
    key (int): 用于加密的密钥，必须是整数。
    message (str): 需要加密的消息字符串。

    返回值:
    str: 加密后的消息字符串。

    示例:
    >>> encrypt_message(4545, 'The affine cipher is a type of monoalphabetic substitution cipher.')
    'VL}p MM{I}p~{HL}Gp{vp pFsH}pxMpyxIx JHL O}F{~pvuOvF{FuF{xIp~{HL}Gi'
    """
    # 将密钥分解为两个部分：keyA 和 keyB
    keyA, keyB = divmod(key, len(SYMBOLS))
    
    # 检查密钥是否有效
    check_keys(keyA, keyB, "encrypt")
    
    # 初始化密文字符串
    cipherText = ""
    
    # 遍历消息中的每个字符
    for symbol in message:
        if symbol in SYMBOLS:
            # 找到符号在符号表中的位置
            symIndex = SYMBOLS.find(symbol)
            # 计算新的符号位置并添加到密文中
            cipherText += SYMBOLS[(symIndex * keyA + keyB) % len(SYMBOLS)]
        else:
            # 如果符号不在符号表中，直接添加到密文中
            cipherText += symbol
    
    return cipherText


def decrypt_message(key: int, message: str) -> str:
    """
    解密消息函数，使用仿射密码算法。
    
    参数:
    key (int): 用于解密的密钥，是一个整数。
    message (str): 需要解密的消息字符串。
    
    返回值:
    str: 解密后的明文字符串。
    
    示例:
    >>> decrypt_message(4545, 'VL}p MM{I}p~{HL}Gp{vp pFsH}pxMpyxIx JHL O}F{~pvuOvF{FuF{xIp~{HL}Gi')
    'The affine cipher is a type of monoalphabetic substitution cipher.'
    """
    # 将密钥分解为两个部分：keyA 和 keyB
    keyA, keyB = divmod(key, len(SYMBOLS))
    
    # 检查密钥是否有效
    check_keys(keyA, keyB, "decrypt")
    
    # 初始化解密后的明文字符串
    plainText = ""
    
    # 计算 keyA 在符号表长度下的模逆元
    modInverseOfkeyA = cryptomath.findModInverse(keyA, len(SYMBOLS))
    
    # 遍历消息中的每个字符
    for symbol in message:
        if symbol in SYMBOLS:
            # 找到字符在符号表中的位置
            symIndex = SYMBOLS.find(symbol)
            # 使用仿射密码公式进行解密，并添加到明文字符串中
            plainText += SYMBOLS[(symIndex - keyB) * modInverseOfkeyA % len(SYMBOLS)]
        else:
            # 如果字符不在符号表中，直接添加到明文字符串中
            plainText += symbol
    
    # 返回解密后的明文字符串
    return plainText


# 定义一个函数，用于生成随机的密钥
def get_random_key():
    # 无限循环，直到找到合适的密钥为止
    while True:
        # 随机生成两个整数 keyA 和 keyB，范围在2到SYMBOLS的长度之间
        keyA = random.randint(2, len(SYMBOLS))
        keyB = random.randint(2, len(SYMBOLS))
        
        # 检查 keyA 和 SYMBOLS 长度的最大公约数是否为1
        # 如果为1，说明 keyA 和 SYMBOLS 长度互质，可以作为密钥的一部分
        if cryptomath.gcd(keyA, len(SYMBOLS)) == 1:
            # 返回计算得到的密钥，由 keyA 和 keyB 组合而成
            return keyA * len(SYMBOLS) + keyB
"""
#例:
def main():
    """
    >>> key = get_random_key()
    >>> msg = "This is a test!"
    >>> decrypt_message(key, encrypt_message(key, msg)) == msg
    True
    """
    message = input("Enter message: ").strip()
    key = int(input("Enter key [2000 - 9000]: ").strip())
    mode = input("Encrypt/Decrypt [E/D]: ").strip().lower()

    if mode.startswith("e"):
        mode = "encrypt"
        translated = encrypt_message(key, message)
    elif mode.startswith("d"):
        mode = "decrypt"
        translated = decrypt_message(key, message)
    print(f"\n{mode.title()}ed text: \n{translated}")

if __name__ == "__main__":
    import doctest

    doctest.testmod()
    main()
"""